#!/usr/local/bin/python
'''
    mrsnapshot.py
    Copyright (C) 2011  Mario Fleischmann

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''

import random
import sys
import commands
from datetime import datetime, timedelta

###############################################################################
'''Config start'''
###############################################################################
zpoolName = "data"
mountPoint = "fileserver"  # leave empty to manage whole zpool
diskPrefix = "ad"  # Prefix for your devices
vdevs = []  # devices in zpool
for i in [4, 6, 8, 10, 12, 14, 16, 18]:
    vdevs.append('/dev/%s%s' % (diskPrefix, i))
dateFormat = "%Y%m%d-%H%M%S"
zfsPath = "/sbin/zfs"
zpoolPath = "/sbin/zpool"
smartctlPath = "/usr/local/sbin/smartctl"
debug = True
verbose = False
testdata = False  # use test data set
really_delete = False  # really delete stuff?
###############################################################################
'''Config end'''
###############################################################################


def nts(number):
    """Convert a single digit number to two-digit string"""
    if len(str(number)) == 1:
        return "0" + str(number)
    else:
        return str(number)

def create_testdata():
    """Generates test data set; requires use of "%Y%m%d-%H%M%S" dateFormat""" 
    print "Generating test data"
    zfsList = ""
    for i in range(2010, 2012):
        for j in range(01, 13):
            if i == 2011 and j > 4:
                break
            for k in range(1, 29):
                for l in range(00, 24):
                    minutes = str(random.randrange(0, 61, 1))
                    seconds = str(random.randrange(0, 61, 1))
                    
                    eightball = random.randrange(0, 60, 1)
                    if eightball == 10: # pick random number for less results
                            # this is ugly and slow as hell
                            zfsList = zfsList + zpoolName + '/' + mountPoint + \
                            '@' + nts(i) + nts(j) + nts(k) + "-" + nts(l) + \
                            nts(minutes) + nts(seconds) + '\n'
    zfsList = zfsList[0:-1]  # cut last empty list element (see var declaration)
    return zfsList            

def createSnapshot():
    """Creates a Snapshot with current time stamp according to dateFormat"""
    if testdata:
        print "Using test data - no snapshot created"
    else:
        snapshotTime = datetime.strftime(now, dateFormat)
        
        if len(mountPoint) == 0:
            snapshotName = "-r %s@%s" % (zpoolName, snapshotTime)
        else:
            snapshotName = "%s/%s@%s" % (zpoolName, mountPoint, snapshotTime)

        command = commands.getstatusoutput('%s snapshot %s' % 
                (zfsPath, snapshotName))
        status, message = command[0], command[1]
        if status != 0:
            print message

def disk_is_standby(dev):
    """Checks if device is in standby mode. Returns boolean"""
    status = commands.getstatusoutput('%s -n standby %s' % (smartctlPath, dev))[0]
    # TODO:    React if disk not found / unavailable / unable to detect device
    #         type etc
    if status == 0:
        return False
    else:
        return True
    
def check_disk_standby():
    """Checks if disks configured in vdev list are in standby. Program will 
    exit if any disk is"""
    disk_standby = False
    for dev in vdevs:
        if disk_is_standby(dev):
            disk_standby = True
            break # no need to test further
    if disk_standby:
        if debug:
            print "At least one disk is in standby"
        sys.exit(1)

def zpool_is_scrubbed(zpool):
    """Checks if scrub is running on zPool. Returns boolean"""
    output = commands.getoutput('%s status %s' % (zpoolPath, zpool))
    if 'scrub in progress' in output:
        return True
    else:
        return False
    
def check_zpool_scrub():
    """Checks if scrub is running on zpool. In this case it will exit as
    creating a snapshots requires restarting the whole scrub"""
    if zpool_is_scrubbed(zpoolName):
        sys.exit(1)
    else:
        if debug:
            "%s is not being scrubbed" % zpoolName

def delete_snapshot(snapshot):
    if debug:
        print "Deleting %s" % snapshot
    if really_delete:
        if "@" in snapshot: # just for safety
            status = commands.getstatusoutput('%s destroy %s' % (zfsPath, snapshot))
            if status != 0:
                print status[1]

def most_recent(snapshots):
    """Returns the most recent of a list of snapshots. Input format:
    [(datetime.datetime), 'snapshotname')]"""
    if snapshots:
        return sorted(snapshots)[-1]
    else:
        return None

def purge(snapshots, delta = timedelta(seconds = 0)):
    """Will delete any snapshot whose time stamp is between the most recent
    snapshot and the set delta. If the snapshot list contains snapshots where
    the interval between two snapshots is larger than delta, the latter will 
    become the most recent.
    CAUTION: If delta == None snapshots will be deleted"""
    if snapshots:    
        start = most_recent(snapshots)[0]
        start_name = most_recent(snapshots)[1][0]
        
        for snapshot in reversed(sorted(snapshots)):
            timestamp = snapshot[0]
            name = snapshot[1][0]
    
            if delta == None: # if no delta is given, show no mercy
                delete_snapshot(name)
                continue
            
            if verbose:
                print "Comparing %s with %s" % (start_name, name)
            
            if start == timestamp:  # these are the same, keep
                if debug:
                    print "Keeping  %s" % name
                continue
            elif start - timestamp >= delta:    # this is larger than our delta
                                                # a new cycle begins
                if debug:
                    print "Keeping  %s" % name
                if verbose:
                    print "New start: " + str(timestamp)
                start, start_name = timestamp, name
            elif start - timestamp < delta:
                if verbose:
                    print str(start - timestamp) + " is smaller than " + str(delta)
                delete_snapshot(name)
            
def zfsList():
    """Queries ZFS for a list of snapshots according to zpoolName and 
    mountPoint (if set)
    Returns dictionary: time stamp is key, [name, size] are values"""
    snapshotDict = {}
    if testdata:
        print "Running in debug mode"
        zfsList = create_testdata()
    else:
        zfsList = commands.getoutput('%s list -H -t snapshot' % zfsPath)
        
    for line in zfsList.splitlines():
        if len(mountPoint) == 0:
            lineStart = "%s@" % zpoolName
        else:
            lineStart = "%s/%s@" % (zpoolName, mountPoint)
            
        if line.startswith("%s" % lineStart):
            splitline = line.split("\t") # split at tabs
            name = splitline[0]
            size = splitline[1] 
            try:                   
                dateSnapshot = datetime.strptime(name.split("@")[1],
                    dateFormat)
                snapshotDict[dateSnapshot] = name, size
            except ValueError:
                continue # ignore snapshots in wrong date format
            except IndexError, e:
                print e
                sys.exit(1)
    return snapshotDict

def main():
    check_disk_standby() # do nothing if disks are in standby
    #check_zpool_scrub() # do nothing if scrub is running
    createSnapshot()
    zfsSnapshots = zfsList() # get snapshots

    current_week = []
    current_month = []
    last_three_months = []
    old = []
    zerosize = []
    
    for snapshot in sorted(zfsSnapshots.iteritems()):
        timestamp = snapshot[0]
        size = snapshot[1][1]
        
        if size == "0":
            zerosize.append(snapshot)
        elif timestamp > yesterday: # this is today
            continue
        elif timestamp > a_week_ago:
            current_week.append(snapshot)
        elif timestamp > a_month_ago:
            current_month.append(snapshot)
        elif timestamp >= three_months_ago:
            last_three_months.append(snapshot)
        elif timestamp < three_months_ago:
            old.append(snapshot)                   
            

    purge(current_week, delta = one_hour)
    purge(current_month, delta = a_day)
    purge(last_three_months, delta = a_week)
    purge(old, delta = a_month)       
    purge(zerosize, delta = None)


# Global vars here:
now = datetime.now()  # the time
# Deltas
one_hour = timedelta(hours = 1)
half_a_day = 12 * one_hour
a_day = 2 * half_a_day
a_week = 7 * a_day
two_weeks = 2 * a_week
three_weeks = 3 * a_week
a_month = 4 * a_week
a_year = 56 * a_week
# Points in time
yesterday = now - a_day
a_week_ago = now - a_week
a_month_ago = now - a_month
three_months_ago = now - a_month * 3

    
main()  # he ho let's go
